package com.zhanghu.backendZray.controller.ai;

import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.zhanghu.backendZray.cache.ContextCache;
import com.zhanghu.backendZray.pojo.entity.AiQuestionRecord;
import com.zhanghu.backendZray.pojo.request.ai.AiRequest;
import com.zhanghu.backendZray.service.AiQuestionRecordService;
import com.zhanghu.backendZray.utils.AI.TongyiUtil;
import io.reactivex.Flowable;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

/**
 * @author ZhangDong
 * @date 2024/10/6 21:26
 */
@RestController
@RequestMapping("ai")
@RequiredArgsConstructor
@Slf4j
public class AiController {
    private final AiQuestionRecordService aiQuestionRecordService;
    @PostMapping("/stream")
    public ResponseEntity<StreamingResponseBody> stream(@RequestBody AiRequest aiRequest) {
        Integer currentUserId = ContextCache.currentUserId();
        StreamingResponseBody responseBody = outputStream -> {
            try {
                Flowable<GenerationResult> flowable = TongyiUtil.streamCallWithMessage(aiRequest.getContent());
                flowable.blockingForEach(message -> {
                    for (GenerationOutput.Choice choice : message.getOutput().getChoices()) {
                        if ("stop".equals(choice.getFinishReason())) {
                            outputStream.write("\n".getBytes());
                        } else {
                            String content = choice.getMessage().getContent();
                            outputStream.write(content.getBytes());
                        }
                        outputStream.flush();
                    }
                });
                AiQuestionRecord record = new AiQuestionRecord();
                record.setUserId(currentUserId);
                record.setQuestion(aiRequest.getContent());
                aiQuestionRecordService.save(record);
            } catch (NoApiKeyException e) {
                throw new RuntimeException(e);
            } catch (InputRequiredException e) {
                throw new RuntimeException(e);
            }
        };
        return ResponseEntity.ok().body(responseBody);
    }
}
